import os
import numpy as np
from sklearn.datasets import load_files
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import train_test_split


class NaiveBayes:
    def __init__(self, alpha=1.0):
        self.alpha = alpha
        self.log_class_priors = None
        self.log_likelihoods = None
        self.classes = None
        self.vocabulary_size = None

    def fit(self, X, y):
        self.classes = np.unique(y)
        num_samples, self.vocabulary_size = X.shape

        # 计算先验概率 P(C) 并取对数
        class_counts = np.bincount(y)
        self.log_class_priors = np.log((class_counts + self.alpha) / (num_samples + self.alpha * len(self.classes)))

        # 计算条件概率 P(X|C)，并应用拉普拉斯平滑
        self.log_likelihoods = np.zeros((len(self.classes), self.vocabulary_size))
        for idx, c in enumerate(self.classes):
            X_c = X[y == c]
            feature_counts = X_c.sum(axis=0) + self.alpha
            self.log_likelihoods[idx, :] = np.log(feature_counts / (X_c.sum() + self.alpha * self.vocabulary_size))

    def predict(self, X):
        scores = X @ self.log_likelihoods.T + self.log_class_priors
        return self.classes[np.argmax(scores, axis=1)]


def load_local_data(local_path, categories=None):
    if not os.path.exists(local_path):
        raise FileNotFoundError(f"The specified path does not exist: {local_path}")

    print("Loading data from local path...")
    newsgroups = load_files(local_path, categories=categories, encoding='latin-1')
    vectorizer = CountVectorizer(binary=True, stop_words='english', max_features=1000)
    X = vectorizer.fit_transform(newsgroups.data).toarray()
    y = newsgroups.target
    return train_test_split(X, y, test_size=0.2, random_state=42)


if __name__ == "__main__":
    try:
        # 指定本地数据集路径
        local_data_path = "D:\\python working\\test001\\20_newsgroups"
        categories = ['sci.space', 'rec.autos']  # 只选择两个类别进行测试

        # 加载数据
        X_train, X_test, y_train, y_test = load_local_data(local_data_path, categories=categories)

        # 初始化并训练模型
        model = NaiveBayes(alpha=1.0)
        model.fit(X_train, y_train)

        # 预测并计算准确率
        y_pred = model.predict(X_test)
        accuracy = np.mean(y_pred == y_test)
        print(f"测试集准确率: {accuracy:.4f}")
    except Exception as e:
        print(f"An error occurred: {e}")